import streamlit as st
from hello_qwen import get_response
from first import latex_to_markdown
import os
import json
from recognize import recognize_image


# 设置页面配置
st.set_page_config(
    page_title="数学助手",
    page_icon="📚",
    layout="wide",
)


# 标题
st.title("数学助手")


def display_chat(messages):
    # 遍历消息列表
    for message in messages:
        role, content = message["role"], message["content"]
        # 根据角色不同，设置不同样式的边框
        if role == "user":
            # 用户消息样式
            border_style = "border-left: 4px solid #ffa07a; margin: 5px 0; padding: 5px; background-color: #fff5f7;"
        else:
            # AI消息样式
            border_style = "border-left: 4px solid #20b2aa; margin: 5px 0; padding: 5px; background-color: #e0ffff;"
        
        # 使用Markdown的HTML实体添加边框样式
        st.markdown(
            f'<div style="{border_style}">'
            f"<strong>{role}:</strong> {content}"
            "</div>",
            unsafe_allow_html=True
        )

def stream_response(response_stream, chat_placeholder):
    global ai_content
    ai_content = ''
    for chunk in response_stream:
        content = chunk.choices[0].delta.content
        if content:  # 确保content非空
            with chat_placeholder:
                ai_content = ai_content + content
                # 使用 Markdown 渲染公式
                st.markdown(f"**AI:** {ai_content}", unsafe_allow_html=True)
                # 显示 LaTeX 公式
                st.markdown(f"**AI (LaTeX):** $${ai_content}$$", unsafe_allow_html=True)

# 主函数
def main():
    print(os.getcwd())
    # 初始化聊天记录
    if "messages" not in st.session_state:
        st.session_state.messages = []

    # 在主区域显示聊天记录
    display_chat(st.session_state.messages)
    # 创建聊天显示区域
    chat_placeholder = st.empty()
    # 用户输入和发送按钮放在页面底部
    user_input = st.text_input("您: ")
    # 输入LaTeX 代码
    latex_url = st.text_input("输入图片地址:")

    # 设置按钮
    send_button = st.button("发送")
    buttonClean = st.sidebar.button("清理会话历史", key="clean")

    # 点击重置按钮时
    if buttonClean:
        st.session_state.messages = []  # 清空聊天记录
        chat_placeholder.empty()  # 清空聊天显示区域

    # 点击发送按钮时
    if send_button and user_input:
        img_path = latex_url  # 图片路径
        inference_mode = 'cpu'  
        num_beam = 1
        use_mix = False  

        latex_code = recognize_image(img_path, inference_mode, num_beam, use_mix)
        # 转换成markdown文本
        markdown_text = latex_to_markdown(user_input, latex_code)
        print(markdown_text)
        response = get_response( markdown_text)
        # print(response)
        stream_response(response, chat_placeholder)
        st.session_state.messages.append({"role": "user", "content": user_input})
        st.session_state.messages.append({"role": "AI", "content": ai_content})
        # # 在侧边栏公式转换结果
        with st.sidebar:
            if user_input:  # 只有当用户输入了内容时才进行搜索
                st.write("公式转换结果：")
                st.write(f"## Formula\n\n$${latex_code}$$\n\n")  # 显示搜索结果


# 运行主函数
if __name__ == "__main__":
    main()

